{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-03-06T16:39:19.190824Z",
     "start_time": "2024-03-06T16:39:18.706794Z"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "data = pd.read_csv(\"E:/baselines/SP500_new.csv\")\n",
    "#data = data.drop(columns=['Adj Close'])\n",
    "#data.rename(columns={'Date':'date', 'Open':'open', 'High':'high', 'Low':'low', 'Close':'close', 'Volume':'volume', 'company':'stockcode'}, inplace=True)\n",
    "data.rename(columns={'dt':'date', 'kdcode':'stockcode'}, inplace=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T16:39:19.776474Z",
     "start_time": "2024-03-06T16:39:19.191918Z"
    }
   },
   "id": "659d9930f6f0659a",
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "       stockcode       date        open        high         low       close  \\\n0              A 2018-01-02   67.419998   67.889999   67.339996   67.599998   \n1              A 2018-01-03   67.620003   69.489998   67.599998   69.320000   \n2              A 2018-01-04   69.540001   69.820000   68.779999   68.800003   \n3              A 2018-01-05   68.730003   70.099998   68.730003   69.900002   \n4              A 2018-01-08   69.730003   70.330002   69.550003   70.050003   \n...          ...        ...         ...         ...         ...         ...   \n734391       ZTS 2023-12-21  196.000000  196.979996  193.899994  194.660004   \n734392       ZTS 2023-12-22  195.320007  195.910004  192.740005  194.979996   \n734393       ZTS 2023-12-26  194.880005  196.339996  194.089996  195.500000   \n734394       ZTS 2023-12-27  195.410004  197.009995  194.740005  196.899994   \n734395       ZTS 2023-12-28  197.619995  198.600006  196.529999  197.160004   \n\n           volume  \n0       1047800.0  \n1       1698900.0  \n2       2230700.0  \n3       1632500.0  \n4       1613400.0  \n...           ...  \n734391  2118600.0  \n734392  1548400.0  \n734393   814600.0  \n734394   766400.0  \n734395   880100.0  \n\n[734396 rows x 7 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>stockcode</th>\n      <th>date</th>\n      <th>open</th>\n      <th>high</th>\n      <th>low</th>\n      <th>close</th>\n      <th>volume</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>A</td>\n      <td>2018-01-02</td>\n      <td>67.419998</td>\n      <td>67.889999</td>\n      <td>67.339996</td>\n      <td>67.599998</td>\n      <td>1047800.0</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>A</td>\n      <td>2018-01-03</td>\n      <td>67.620003</td>\n      <td>69.489998</td>\n      <td>67.599998</td>\n      <td>69.320000</td>\n      <td>1698900.0</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>A</td>\n      <td>2018-01-04</td>\n      <td>69.540001</td>\n      <td>69.820000</td>\n      <td>68.779999</td>\n      <td>68.800003</td>\n      <td>2230700.0</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>A</td>\n      <td>2018-01-05</td>\n      <td>68.730003</td>\n      <td>70.099998</td>\n      <td>68.730003</td>\n      <td>69.900002</td>\n      <td>1632500.0</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>A</td>\n      <td>2018-01-08</td>\n      <td>69.730003</td>\n      <td>70.330002</td>\n      <td>69.550003</td>\n      <td>70.050003</td>\n      <td>1613400.0</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>734391</th>\n      <td>ZTS</td>\n      <td>2023-12-21</td>\n      <td>196.000000</td>\n      <td>196.979996</td>\n      <td>193.899994</td>\n      <td>194.660004</td>\n      <td>2118600.0</td>\n    </tr>\n    <tr>\n      <th>734392</th>\n      <td>ZTS</td>\n      <td>2023-12-22</td>\n      <td>195.320007</td>\n      <td>195.910004</td>\n      <td>192.740005</td>\n      <td>194.979996</td>\n      <td>1548400.0</td>\n    </tr>\n    <tr>\n      <th>734393</th>\n      <td>ZTS</td>\n      <td>2023-12-26</td>\n      <td>194.880005</td>\n      <td>196.339996</td>\n      <td>194.089996</td>\n      <td>195.500000</td>\n      <td>814600.0</td>\n    </tr>\n    <tr>\n      <th>734394</th>\n      <td>ZTS</td>\n      <td>2023-12-27</td>\n      <td>195.410004</td>\n      <td>197.009995</td>\n      <td>194.740005</td>\n      <td>196.899994</td>\n      <td>766400.0</td>\n    </tr>\n    <tr>\n      <th>734395</th>\n      <td>ZTS</td>\n      <td>2023-12-28</td>\n      <td>197.619995</td>\n      <td>198.600006</td>\n      <td>196.529999</td>\n      <td>197.160004</td>\n      <td>880100.0</td>\n    </tr>\n  </tbody>\n</table>\n<p>734396 rows × 7 columns</p>\n</div>"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['date'] = pd.to_datetime(data['date'])\n",
    "grouped_df = data.groupby('stockcode')\n",
    "l = data['date'].nunique()\n",
    "def filter_groups(group):\n",
    "    return len(group) == l\n",
    "df = grouped_df.filter(filter_groups)\n",
    "df"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T16:39:20.165022Z",
     "start_time": "2024-03-06T16:39:19.777783Z"
    }
   },
   "id": "c1a98a0060e8115a",
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "487"
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['stockcode'].nunique()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T16:39:20.227630Z",
     "start_time": "2024-03-06T16:39:20.167031Z"
    }
   },
   "id": "7f3f4ae5492b2e10",
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "       stockcode       date        open        high         low       close  \\\n0              A 2018-01-02   67.419998   67.889999   67.339996   67.599998   \n1              A 2018-01-03   67.620003   69.489998   67.599998   69.320000   \n2              A 2018-01-04   69.540001   69.820000   68.779999   68.800003   \n3              A 2018-01-05   68.730003   70.099998   68.730003   69.900002   \n4              A 2018-01-08   69.730003   70.330002   69.550003   70.050003   \n...          ...        ...         ...         ...         ...         ...   \n734391       ZTS 2023-12-21  196.000000  196.979996  193.899994  194.660004   \n734392       ZTS 2023-12-22  195.320007  195.910004  192.740005  194.979996   \n734393       ZTS 2023-12-26  194.880005  196.339996  194.089996  195.500000   \n734394       ZTS 2023-12-27  195.410004  197.009995  194.740005  196.899994   \n734395       ZTS 2023-12-28  197.619995  198.600006  196.529999  197.160004   \n\n           volume     ret_1     label  \n0       1047800.0  1.000000  1.025444  \n1       1698900.0  1.025444  0.992499  \n2       2230700.0  0.992499  1.015988  \n3       1632500.0  1.015988  1.002146  \n4       1613400.0  1.002146  1.024554  \n...           ...       ...       ...  \n734391  2118600.0  1.000154  1.001644  \n734392  1548400.0  1.001644  1.002667  \n734393   814600.0  1.002667  1.007161  \n734394   766400.0  1.007161  1.001321  \n734395   880100.0  1.001321  1.000000  \n\n[734396 rows x 9 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>stockcode</th>\n      <th>date</th>\n      <th>open</th>\n      <th>high</th>\n      <th>low</th>\n      <th>close</th>\n      <th>volume</th>\n      <th>ret_1</th>\n      <th>label</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>A</td>\n      <td>2018-01-02</td>\n      <td>67.419998</td>\n      <td>67.889999</td>\n      <td>67.339996</td>\n      <td>67.599998</td>\n      <td>1047800.0</td>\n      <td>1.000000</td>\n      <td>1.025444</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>A</td>\n      <td>2018-01-03</td>\n      <td>67.620003</td>\n      <td>69.489998</td>\n      <td>67.599998</td>\n      <td>69.320000</td>\n      <td>1698900.0</td>\n      <td>1.025444</td>\n      <td>0.992499</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>A</td>\n      <td>2018-01-04</td>\n      <td>69.540001</td>\n      <td>69.820000</td>\n      <td>68.779999</td>\n      <td>68.800003</td>\n      <td>2230700.0</td>\n      <td>0.992499</td>\n      <td>1.015988</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>A</td>\n      <td>2018-01-05</td>\n      <td>68.730003</td>\n      <td>70.099998</td>\n      <td>68.730003</td>\n      <td>69.900002</td>\n      <td>1632500.0</td>\n      <td>1.015988</td>\n      <td>1.002146</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>A</td>\n      <td>2018-01-08</td>\n      <td>69.730003</td>\n      <td>70.330002</td>\n      <td>69.550003</td>\n      <td>70.050003</td>\n      <td>1613400.0</td>\n      <td>1.002146</td>\n      <td>1.024554</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>734391</th>\n      <td>ZTS</td>\n      <td>2023-12-21</td>\n      <td>196.000000</td>\n      <td>196.979996</td>\n      <td>193.899994</td>\n      <td>194.660004</td>\n      <td>2118600.0</td>\n      <td>1.000154</td>\n      <td>1.001644</td>\n    </tr>\n    <tr>\n      <th>734392</th>\n      <td>ZTS</td>\n      <td>2023-12-22</td>\n      <td>195.320007</td>\n      <td>195.910004</td>\n      <td>192.740005</td>\n      <td>194.979996</td>\n      <td>1548400.0</td>\n      <td>1.001644</td>\n      <td>1.002667</td>\n    </tr>\n    <tr>\n      <th>734393</th>\n      <td>ZTS</td>\n      <td>2023-12-26</td>\n      <td>194.880005</td>\n      <td>196.339996</td>\n      <td>194.089996</td>\n      <td>195.500000</td>\n      <td>814600.0</td>\n      <td>1.002667</td>\n      <td>1.007161</td>\n    </tr>\n    <tr>\n      <th>734394</th>\n      <td>ZTS</td>\n      <td>2023-12-27</td>\n      <td>195.410004</td>\n      <td>197.009995</td>\n      <td>194.740005</td>\n      <td>196.899994</td>\n      <td>766400.0</td>\n      <td>1.007161</td>\n      <td>1.001321</td>\n    </tr>\n    <tr>\n      <th>734395</th>\n      <td>ZTS</td>\n      <td>2023-12-28</td>\n      <td>197.619995</td>\n      <td>198.600006</td>\n      <td>196.529999</td>\n      <td>197.160004</td>\n      <td>880100.0</td>\n      <td>1.001321</td>\n      <td>1.000000</td>\n    </tr>\n  </tbody>\n</table>\n<p>734396 rows × 9 columns</p>\n</div>"
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['ret_1'] = df['close'] / df['close'].shift(1)\n",
    "df['ret_1'] = df['ret_1'].fillna(1)\n",
    "df['label'] = df['ret_1'].shift(-1)\n",
    "df['label'] = df['label'].fillna(1)\n",
    "df"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T16:39:20.275226Z",
     "start_time": "2024-03-06T16:39:20.230697Z"
    }
   },
   "id": "f45638e23043b6fd",
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "df.loc[df['date'] == '2018-01-02', 'ret_1'] = 1\n",
    "df.loc[df['date'] == '2023-12-28', 'label'] = 1"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T16:39:20.291115Z",
     "start_time": "2024-03-06T16:39:20.276417Z"
    }
   },
   "id": "aa2fadab764493a5",
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "       stockcode       date        open        high         low       close  \\\n0              A 2018-01-02   67.419998   67.889999   67.339996   67.599998   \n1              A 2018-01-03   67.620003   69.489998   67.599998   69.320000   \n2              A 2018-01-04   69.540001   69.820000   68.779999   68.800003   \n3              A 2018-01-05   68.730003   70.099998   68.730003   69.900002   \n4              A 2018-01-08   69.730003   70.330002   69.550003   70.050003   \n...          ...        ...         ...         ...         ...         ...   \n734391       ZTS 2023-12-21  196.000000  196.979996  193.899994  194.660004   \n734392       ZTS 2023-12-22  195.320007  195.910004  192.740005  194.979996   \n734393       ZTS 2023-12-26  194.880005  196.339996  194.089996  195.500000   \n734394       ZTS 2023-12-27  195.410004  197.009995  194.740005  196.899994   \n734395       ZTS 2023-12-28  197.619995  198.600006  196.529999  197.160004   \n\n           volume     ret_1     label  \n0       1047800.0  1.000000  1.025444  \n1       1698900.0  1.025444  0.992499  \n2       2230700.0  0.992499  1.015988  \n3       1632500.0  1.015988  1.002146  \n4       1613400.0  1.002146  1.024554  \n...           ...       ...       ...  \n734391  2118600.0  1.000154  1.001644  \n734392  1548400.0  1.001644  1.002667  \n734393   814600.0  1.002667  1.007161  \n734394   766400.0  1.007161  1.001321  \n734395   880100.0  1.001321  1.000000  \n\n[734396 rows x 9 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>stockcode</th>\n      <th>date</th>\n      <th>open</th>\n      <th>high</th>\n      <th>low</th>\n      <th>close</th>\n      <th>volume</th>\n      <th>ret_1</th>\n      <th>label</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>A</td>\n      <td>2018-01-02</td>\n      <td>67.419998</td>\n      <td>67.889999</td>\n      <td>67.339996</td>\n      <td>67.599998</td>\n      <td>1047800.0</td>\n      <td>1.000000</td>\n      <td>1.025444</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>A</td>\n      <td>2018-01-03</td>\n      <td>67.620003</td>\n      <td>69.489998</td>\n      <td>67.599998</td>\n      <td>69.320000</td>\n      <td>1698900.0</td>\n      <td>1.025444</td>\n      <td>0.992499</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>A</td>\n      <td>2018-01-04</td>\n      <td>69.540001</td>\n      <td>69.820000</td>\n      <td>68.779999</td>\n      <td>68.800003</td>\n      <td>2230700.0</td>\n      <td>0.992499</td>\n      <td>1.015988</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>A</td>\n      <td>2018-01-05</td>\n      <td>68.730003</td>\n      <td>70.099998</td>\n      <td>68.730003</td>\n      <td>69.900002</td>\n      <td>1632500.0</td>\n      <td>1.015988</td>\n      <td>1.002146</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>A</td>\n      <td>2018-01-08</td>\n      <td>69.730003</td>\n      <td>70.330002</td>\n      <td>69.550003</td>\n      <td>70.050003</td>\n      <td>1613400.0</td>\n      <td>1.002146</td>\n      <td>1.024554</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>734391</th>\n      <td>ZTS</td>\n      <td>2023-12-21</td>\n      <td>196.000000</td>\n      <td>196.979996</td>\n      <td>193.899994</td>\n      <td>194.660004</td>\n      <td>2118600.0</td>\n      <td>1.000154</td>\n      <td>1.001644</td>\n    </tr>\n    <tr>\n      <th>734392</th>\n      <td>ZTS</td>\n      <td>2023-12-22</td>\n      <td>195.320007</td>\n      <td>195.910004</td>\n      <td>192.740005</td>\n      <td>194.979996</td>\n      <td>1548400.0</td>\n      <td>1.001644</td>\n      <td>1.002667</td>\n    </tr>\n    <tr>\n      <th>734393</th>\n      <td>ZTS</td>\n      <td>2023-12-26</td>\n      <td>194.880005</td>\n      <td>196.339996</td>\n      <td>194.089996</td>\n      <td>195.500000</td>\n      <td>814600.0</td>\n      <td>1.002667</td>\n      <td>1.007161</td>\n    </tr>\n    <tr>\n      <th>734394</th>\n      <td>ZTS</td>\n      <td>2023-12-27</td>\n      <td>195.410004</td>\n      <td>197.009995</td>\n      <td>194.740005</td>\n      <td>196.899994</td>\n      <td>766400.0</td>\n      <td>1.007161</td>\n      <td>1.001321</td>\n    </tr>\n    <tr>\n      <th>734395</th>\n      <td>ZTS</td>\n      <td>2023-12-28</td>\n      <td>197.619995</td>\n      <td>198.600006</td>\n      <td>196.529999</td>\n      <td>197.160004</td>\n      <td>880100.0</td>\n      <td>1.001321</td>\n      <td>1.000000</td>\n    </tr>\n  </tbody>\n</table>\n<p>734396 rows × 9 columns</p>\n</div>"
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T16:39:20.306433Z",
     "start_time": "2024-03-06T16:39:20.292628Z"
    }
   },
   "id": "9a6612265a9052d0",
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 487/487 [00:19<00:00, 25.45it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "scaler = MinMaxScaler()\n",
    "cols_to_normalize = ['close', 'open', 'high', 'low', 'volume']\n",
    "# unique_stockcodes = df['stockcode'].unique()[:100]\n",
    "\n",
    "# df = df[df['stockcode'].isin(unique_stockcodes)]\n",
    "\n",
    "for stock in tqdm(df['stockcode'].unique()):\n",
    "    tmp = df[df['stockcode'] == stock].reset_index(drop=True)\n",
    "    tmp[cols_to_normalize] = scaler.fit_transform(tmp[cols_to_normalize])\n",
    "    tmp.to_pickle(f\"sp500/{stock}.pkl\")"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T16:40:47.178505Z",
     "start_time": "2024-03-06T16:40:28.014384Z"
    }
   },
   "id": "748416f223963314",
   "execution_count": 10
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "df.to_csv('sp500_fea.csv', index=None)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-06T17:20:19.885053Z",
     "start_time": "2024-03-06T17:20:14.785094Z"
    }
   },
   "id": "e6d5f7f825b9f0f0",
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "df['stockcode'].unique()"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ab628673d8222818",
   "execution_count": null
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "pom = pd.read_pickle(\"zz500/000009.SZ.pkl\")\n",
    "pom"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "start_time": "2024-03-06T16:39:25.526580Z"
    }
   },
   "id": "673880f900d81fb3",
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
